import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# -------------------------- 解决中文显示问题 --------------------------
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]  # 支持中文的字体
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# -------------------------- 1. 定义核心参数 --------------------------
# 假目标（原点）
O = np.array([0, 0, 0])

# 真目标（圆柱：底面圆心T，半径r=7m，高h=10m）
T = np.array([0, 200, 0])  # 底面圆心（z=0）
r_true = 7
h_true = 10

# 导弹初始位置与轨迹方向向量（方向指向假目标O）
missiles = {
    'M1': {'pos0': np.array([20000, 0, 2000]), 'color': 'red'},
    'M2': {'pos0': np.array([19000, 600, 2100]), 'color': 'orange'},
    'M3': {'pos0': np.array([18000, -600, 1900]), 'color': 'darkred'}
}
v_missile = 300  # 导弹速度，单位m/s

# 无人机初始位置
uavs = {
    'FY1': {'pos0': np.array([17800, 0, 1800]), 'color': 'blue'},
    'FY2': {'pos0': np.array([12000, 1400, 1400]), 'color': 'green'},
    'FY3': {'pos0': np.array([6000, -3000, 700]), 'color': 'purple'},
    'FY4': {'pos0': np.array([11000, 2000, 1800]), 'color': 'cyan'},
    'FY5': {'pos0': np.array([13000, -2000, 1300]), 'color': 'magenta'}
}

# -------------------------- 2. 初始化3D图 --------------------------
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 设置坐标轴范围（聚焦关键区域）
ax.set_xlim([17000, 21000])
ax.set_ylim([-700, 700])
ax.set_zlim([0, 2200])
ax.set_xlabel('X轴（m）', fontsize=12)
ax.set_ylabel('Y轴（m）', fontsize=12)
ax.set_zlabel('Z轴（m）', fontsize=12)
ax.set_title('烟幕干扰弹投放场景空间可视化\n（假目标O为原点，真目标T为圆柱，导弹/无人机为初始位置+轨迹）', fontsize=14)
ax.grid(True, alpha=0.3)

# -------------------------- 3. 绘制假目标 --------------------------
ax.scatter(O[0], O[1], O[2], s=200, c='black', marker='*', label='假目标O（原点）')

# -------------------------- 4. 绘制真目标（圆柱） --------------------------
theta = np.linspace(0, 2*np.pi, 36)  # 圆的角度参数
# 底面（z=0）顶点
bottom_circle = np.vstack([
    T[0] + r_true * np.cos(theta),
    T[1] + r_true * np.sin(theta),
    np.zeros_like(theta)
]).T
# 顶面（z=10）顶点
top_circle = np.vstack([
    T[0] + r_true * np.cos(theta),
    T[1] + r_true * np.sin(theta),
    np.full_like(theta, h_true)
]).T
# 绘制上下底面
ax.add_collection3d(Poly3DCollection([bottom_circle], facecolors='lightgreen', alpha=0.5, label='真目标（底面）'))
ax.add_collection3d(Poly3DCollection([top_circle], facecolors='lightgreen', alpha=0.5))
# 绘制圆柱侧面
for i in range(len(theta)):
    side = np.vstack([bottom_circle[i], top_circle[i]])
    ax.plot(side[:, 0], side[:, 1], side[:, 2], c='green', alpha=0.5)

# -------------------------- 5. 绘制导弹（初始位置+飞行轨迹） --------------------------
traj_length = 2000  # 轨迹延长长度，单位m
for name, data in missiles.items():
    pos0 = data['pos0']
    color = data['color']
    # 计算方向向量
    dir_vec = O - pos0
    dir_unit = dir_vec / np.linalg.norm(dir_vec)
    traj_end = pos0 + dir_unit * traj_length
    # 绘制初始位置
    ax.scatter(pos0[0], pos0[1], pos0[2], s=150, c=color, marker='^', label=f'导弹{name}（初始位置）')
    # 绘制轨迹
    ax.plot([pos0[0], traj_end[0]], [pos0[1], traj_end[1]], [pos0[2], traj_end[2]], c=color, alpha=0.7, linewidth=2)

# -------------------------- 6. 绘制无人机（初始位置） --------------------------
for name, data in uavs.items():
    pos0 = data['pos0']
    color = data['color']
    ax.scatter(pos0[0], pos0[1], pos0[2], s=120, c=color, marker='s', label=f'无人机{name}（初始位置）')

# -------------------------- 7. 调整图例与视角 --------------------------
ax.legend(loc='upper right', bbox_to_anchor=(1.3, 1.0), fontsize=10)
ax.view_init(elev=20, azim=45)

# 显示图像
plt.tight_layout()
plt.show()
